#!/bin/bash

# SSH Slurm MCP Setup Script for Claude Desktop
# This script automates the setup process

set -e  # Exit on any error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MCP_DIR="$SCRIPT_DIR"

echo -e "${BLUE}🚀 SSH Slurm MCP Setup for Claude Desktop${NC}"
echo -e "${BLUE}================================================${NC}"
echo ""

# Step 1: Check if we're in the right directory
echo -e "${YELLOW}📁 Checking MCP directory...${NC}"
if [[ ! -f "$MCP_DIR/server.py" ]]; then
    echo -e "${RED}❌ Error: server.py not found in $MCP_DIR${NC}"
    echo "Please run this script from the ssh-slurm-mcp directory"
    exit 1
fi
echo -e "${GREEN}✅ Found MCP server files${NC}"

# Step 2: Make run script executable
echo -e "${YELLOW}🔧 Setting up permissions...${NC}"
chmod +x "$MCP_DIR/run.sh"
echo -e "${GREEN}✅ Made run.sh executable${NC}"

# Step 3: Check Python and install dependencies
echo -e "${YELLOW}🐍 Checking Python dependencies...${NC}"

# Check if uv is available
if command -v uv &> /dev/null; then
    echo "Using uv package manager..."
    cd "$MCP_DIR"
    uv sync
    echo -e "${GREEN}✅ Dependencies installed with uv${NC}"
elif command -v pip &> /dev/null; then
    echo "Using pip package manager..."
    if [[ -f "$MCP_DIR/requirements.txt" ]]; then
        pip install -r "$MCP_DIR/requirements.txt"
        echo -e "${GREEN}✅ Dependencies installed with pip${NC}"
    else
        echo "Installing core dependencies..."
        pip install fastmcp paramiko asyncio pathlib
        echo -e "${GREEN}✅ Core dependencies installed${NC}"
    fi
else
    echo -e "${RED}❌ Error: Neither uv nor pip found${NC}"
    echo "Please install Python package manager first"
    exit 1
fi

# Step 4: Test the MCP server
echo -e "${YELLOW}🧪 Testing MCP server...${NC}"
cd "$MCP_DIR"

# Test server startup (timeout after 5 seconds)
if timeout 5 python server.py > /dev/null 2>&1; then
    echo -e "${GREEN}✅ MCP server test completed${NC}"
else
    # Check if it failed due to timeout (expected) or actual error
    if python -c "import sys; sys.path.insert(0, 'src'); from server import mcp; print('Import test passed')" 2>/dev/null; then
        echo -e "${GREEN}✅ MCP server validation passed${NC}"
    else
        echo -e "${RED}❌ MCP server test failed${NC}"
        echo "Please check the server.py file and dependencies"
        exit 1
    fi
fi

# Step 5: Add to Claude Desktop
echo -e "${YELLOW}🔧 Adding to Claude Desktop...${NC}"

# Check if claude CLI is available
if command -v claude &> /dev/null; then
    echo "Using Claude CLI..."
    claude mcp add ssh-slurm-mcp python "$MCP_DIR/run.sh"
    echo -e "${GREEN}✅ Added ssh-slurm-mcp to Claude Desktop using CLI${NC}"
    
    echo ""
    echo -e "${BLUE}📋 Current MCP servers:${NC}"
    claude mcp list
else
    echo -e "${YELLOW}⚠️  Claude CLI not found. Providing manual configuration...${NC}"
    echo ""
    echo -e "${BLUE}📝 Manual Configuration Instructions:${NC}"
    echo "Add this to your Claude Desktop configuration file:"
    echo ""
    echo -e "${GREEN}{"
    echo '  "mcpServers": {'
    echo '    "ssh-slurm-mcp": {'
    echo "      \"command\": \"$MCP_DIR/run.sh\","
    echo '      "args": []'
    echo '    }'
    echo '  }'
    echo -e "}${NC}"
    echo ""
    echo -e "${BLUE}📍 Configuration file locations:${NC}"
    echo "• macOS: ~/Library/Application Support/Claude/claude_desktop_config.json"
    echo "• Windows: %APPDATA%\\Claude\\claude_desktop_config.json"  
    echo "• Linux: ~/.config/claude/claude_desktop_config.json"
fi

# Step 6: Final instructions
echo ""
echo -e "${GREEN}🎉 Setup Complete!${NC}"
echo ""
echo -e "${BLUE}📋 Next Steps:${NC}"
echo "1. 🔄 Restart Claude Desktop application"
echo "2. 🗣️  Start a new conversation in Claude Desktop"
echo "3. 🧪 Test the MCP tools with a query like:"
echo "   ${GREEN}'Help me connect to my SSH server and run some commands'${NC}"
echo ""
echo -e "${BLUE}🛠️  Available Tools:${NC}"
echo "• connect_ssh - Connect to SSH servers"
echo "• execute_sequential_commands - Run multiple commands in sequence"
echo "• compile_code - Compile C/C++/Python code (gcc, hipcc, python)"
echo "• submit_slurm_job - Submit Slurm batch jobs"
echo "• upload_file / download_file - Transfer files"
echo "• And 8 more tools for complete HPC workflow management"
echo ""
echo -e "${BLUE}📚 Documentation:${NC}"
echo "• Setup guide: $MCP_DIR/CLAUDE_SETUP.md"
echo "• Usage examples: $MCP_DIR/examples/"
echo "• Full documentation: $MCP_DIR/README.md"
echo ""
echo -e "${YELLOW}💡 Tip: Ask Claude to list MCP tools to verify the installation!${NC}"
